BRANCH-MoE: Balance-Aware Tree Routing for Large Embedding Models
BRANCH-MoE：面向大规模嵌入模型的平衡感知树路由
🧭 MoE 树形路由与负载均衡 ★ 深度解读
Google Research · University of Southern California
提出将专家置于二叉树叶节点的BRANCH-MoE，用到达加权EMA锚点在没有辅助损失的情况下保持负载均衡并诱导共路由局部性
📌 工作参考价值（4 点）
- 在大规模 CTR/推荐模型中用 MoE 做稀疏专家时，可把 flat top-K 门控换成树形路由：每个内部节点用到达加权均值做 score 中心化（训练用 batch 均值，推理用 EMA），可省掉 auxiliary load-balancing loss 及其调参，同时降低子树饿死风险。
- 树路径天然给出专家二进制地址，可按 prefix 把专家/embedding 分配到设备或参数服务器域；推理时可依据上层 margin 预判是否跨域通信，降低 token 分发成本，适合多级缓存/分片的大规模稀疏 embedding 部署。
- 工程实现上：线性 node map 的树结构局部性最强（UCI 上树距离 0.71 vs 随机 0.817）；换成 MLP 会削弱局部性并提高 load CV，需在目标数据上验证。对输入特征要避免稀有二值列被 z-score 放大成巨大 spike，否则会诱导专家 collapse（Covertype 上需 pin 低方差列）。
- 该方法无辅助损失即可均衡负载，但 Criteo 上 load CV 高于 flat 方法（0.115 vs ≤0.006）；若业务对 per-batch 均衡极敏感，可能仍需要 capacity factor 或调度兜底。